"""
STORK Sampler Nodes for ComfyUI
"""

import comfy.samplers
from .stork_sampler import sample_stork2, sample_stork4

class STORKSamplerSelect:
    @classmethod
    def INPUT_TYPES(s):
        return {"required":
                    {"solver_order": ([2, 4], {"default": 2}),
                     "substeps": ("INT", {"default": 15, "min": 10, "max": 200, "step": 5}),
                     "derivative_order": ([2, 3], {"default": 2}),
                     "debug": ("BOOLEAN", {"default": False}),
                    }
               }
    RETURN_TYPES = ("SAMPLER",)
    CATEGORY = "sampling/42lux/stork"

    FUNCTION = "get_sampler"

    def get_sampler(self, solver_order, substeps, derivative_order, debug):
        if solver_order == 2:
            sampler = comfy.samplers.KSAMPLER(
                lambda *args, **kwargs: sample_stork2(*args, **kwargs, s=substeps, derivative_order=derivative_order, debug=debug)
            )
        else:  # solver_order == 4
            sampler = comfy.samplers.KSAMPLER(
                lambda *args, **kwargs: sample_stork4(*args, **kwargs, s=substeps, derivative_order=derivative_order, debug=debug)
            )
        return (sampler,)

def register_stork_samplers():
    """Register STORK samplers with ComfyUI's sampler system"""
    if "stork2" not in comfy.samplers.SAMPLER_NAMES:
        comfy.samplers.SAMPLER_NAMES.append("stork2")
    if "stork4" not in comfy.samplers.SAMPLER_NAMES:
        comfy.samplers.SAMPLER_NAMES.append("stork4")
    
    if hasattr(comfy.samplers.KSampler, 'SAMPLERS'):
        if "stork2" not in comfy.samplers.KSampler.SAMPLERS:
            comfy.samplers.KSampler.SAMPLERS.append("stork2")
        if "stork4" not in comfy.samplers.KSampler.SAMPLERS:
            comfy.samplers.KSampler.SAMPLERS.append("stork4")

def sampler_object_stork(name):
    """Create STORK sampler object for ComfyUI"""
    if name == "stork2":
        return comfy.samplers.KSAMPLER(sample_stork2)
    elif name == "stork4":
        return comfy.samplers.KSAMPLER(sample_stork4)
    return None

original_sampler_object = comfy.samplers.sampler_object

def patched_sampler_object(name):
    stork_sampler = sampler_object_stork(name)
    if stork_sampler is not None:
        return stork_sampler
    return original_sampler_object(name)

comfy.samplers.sampler_object = patched_sampler_object

register_stork_samplers()

NODE_CLASS_MAPPINGS = {
    "STORKSamplerSelect": STORKSamplerSelect,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "STORKSamplerSelect": "STORK Sampler (42lux)",
}

print("42lux: STORK sampler nodes loaded successfully!")
print("42lux: Available samplers: stork2, stork4")